!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief Module with utility to perform MD Nudged Elastic Band Calculation
!> \note
!>      Numerical accuracy for parallel runs:
!>       Each replica starts the SCF run from the one optimized
!>       in a previous run. It may happen then energies and derivatives
!>       of a serial run and a parallel run could be slightly different
!>       'cause of a different starting density matrix.
!>       Exact results are obtained using:
!>          EXTRAPOLATION USE_GUESS in QS section (Teo 09.2006)
!> \author Teodoro Laino 10.2006
! **************************************************************************************************
MODULE neb_md_utils
   USE cp_units,                        ONLY: cp_unit_from_cp2k
   USE global_types,                    ONLY: global_environment_type
   USE input_constants,                 ONLY: band_md_opt,&
                                              do_band_collective
   USE input_section_types,             ONLY: section_vals_get,&
                                              section_vals_get_subs_vals,&
                                              section_vals_type,&
                                              section_vals_val_get
   USE kinds,                           ONLY: dp
   USE neb_types,                       ONLY: neb_type,&
                                              neb_var_type
   USE particle_types,                  ONLY: get_particle_pos_or_vel,&
                                              particle_type,&
                                              update_particle_pos_or_vel
   USE physcon,                         ONLY: kelvin,&
                                              massunit
#include "../base/base_uses.f90"

   IMPLICIT NONE
   PRIVATE
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'neb_md_utils'

   PUBLIC :: neb_initialize_velocity, &
             control_vels_a, &
             control_vels_b, &
             get_temperatures

CONTAINS

! **************************************************************************************************
!> \brief Initialize velocities of replica in an MD optimized algorithm within
!>        NEB
!> \param vels ...
!> \param neb_section ...
!> \param particle_set ...
!> \param i_rep ...
!> \param iw ...
!> \param globenv ...
!> \param neb_env ...
!> \par History
!>      25.11.2010 Consider core-shell model (MK)
!> \author Teodoro Laino 09.2006
! **************************************************************************************************
   SUBROUTINE neb_initialize_velocity(vels, neb_section, particle_set, i_rep, iw, &
                                      globenv, neb_env)

      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: vels
      TYPE(section_vals_type), POINTER                   :: neb_section
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      INTEGER, INTENT(IN)                                :: i_rep, iw
      TYPE(global_environment_type), POINTER             :: globenv
      TYPE(neb_type), POINTER                            :: neb_env

      INTEGER                                            :: iatom, ivar, natom, nparticle, nvar
      REAL(KIND=dp)                                      :: akin, mass, mass_tot, sc, temp, &
                                                            temp_ext, tmp_r1
      REAL(KIND=dp), DIMENSION(3)                        :: v, vcom
      TYPE(section_vals_type), POINTER                   :: md_section

      IF (neb_env%opt_type == band_md_opt) THEN
         md_section => section_vals_get_subs_vals(neb_section, "OPTIMIZE_BAND%MD")
         CALL section_vals_val_get(md_section, "TEMPERATURE", r_val=temp_ext)
         ! Initialize velocity according to external temperature
         nparticle = SIZE(vels, 1)
         natom = SIZE(particle_set)
         mass_tot = 0.0_dp
         vcom(1:3) = 0.0_dp
         CALL globenv%gaussian_rng_stream%fill(vels(:, i_rep))
         ! Check always if BAND is working in Cartesian or in internal coordinates
         ! If working in cartesian coordinates let's get rid of the COM
         ! Compute also the total mass (both in Cartesian and internal)
         IF (neb_env%use_colvar) THEN
            nvar = nparticle
            mass_tot = REAL(nvar, KIND=dp)*massunit
         ELSE
            DO iatom = 1, natom
               mass = particle_set(iatom)%atomic_kind%mass
               mass_tot = mass_tot + mass
               v(1:3) = get_particle_pos_or_vel(iatom, particle_set, vels(:, i_rep))
               vcom(1:3) = vcom(1:3) + mass*v(1:3)
            END DO
            vcom(1:3) = vcom(1:3)/mass_tot
         END IF
         ! Compute kinetic energy and temperature
         akin = 0.0_dp
         IF (neb_env%use_colvar) THEN
            DO ivar = 1, nvar
               akin = akin + 0.5_dp*massunit*vels(ivar, i_rep)*vels(ivar, i_rep)
            END DO
         ELSE
            DO iatom = 1, natom
               mass = particle_set(iatom)%atomic_kind%mass
               v(1:3) = -vcom(1:3)
               CALL update_particle_pos_or_vel(iatom, particle_set, v(1:3), vels(:, i_rep))
               akin = akin + 0.5_dp*mass*DOT_PRODUCT(v(1:3), v(1:3))
            END DO
            nvar = 3*natom
         END IF
         temp = 2.0_dp*akin/REAL(nvar, KIND=dp)
         ! Scale velocities to get the correct initial temperature and
         sc = SQRT(temp_ext/temp)
         vels(:, i_rep) = vels(:, i_rep)*sc
         ! Re-compute kinetic energya and temperature and velocity of COM
         akin = 0.0_dp
         vcom = 0.0_dp
         IF (neb_env%use_colvar) THEN
            DO ivar = 1, nvar
               akin = akin + 0.5_dp*massunit*vels(ivar, i_rep)*vels(ivar, i_rep)
            END DO
         ELSE
            DO iatom = 1, natom
               mass = particle_set(iatom)%atomic_kind%mass
               v(1:3) = get_particle_pos_or_vel(iatom, particle_set, vels(:, i_rep))
               vcom(1:3) = vcom(1:3) + mass*v(1:3)
               akin = akin + 0.5_dp*mass*DOT_PRODUCT(v(1:3), v(1:3))
            END DO
         END IF
         vcom(1:3) = vcom(1:3)/mass_tot
         ! Dump information
         IF (iw > 0) THEN
            temp = 2.0_dp*akin/REAL(nvar, KIND=dp)
            tmp_r1 = cp_unit_from_cp2k(temp, "K")
            WRITE (iw, '(A,T61,F18.2,A2)') &
               ' NEB| Initial Temperature ', tmp_r1, " K"
            WRITE (iw, '(A,T61,F20.12)') &
               ' NEB| Centre of mass velocity in direction x:', vcom(1), &
               ' NEB| Centre of mass velocity in direction y:', vcom(2), &
               ' NEB| Centre of mass velocity in direction z:', vcom(3)
            WRITE (iw, '(T2,"NEB|",75("*"))')
         END IF
      ELSE
         vels(:, i_rep) = 0.0_dp
      END IF

   END SUBROUTINE neb_initialize_velocity

! **************************************************************************************************
!> \brief Control on  velocities - I part
!> \param vels ...
!> \param particle_set ...
!> \param tc_section ...
!> \param vc_section ...
!> \param output_unit ...
!> \param istep ...
!> \author Teodoro Laino 09.2006
! **************************************************************************************************
   SUBROUTINE control_vels_a(vels, particle_set, tc_section, vc_section, &
                             output_unit, istep)
      TYPE(neb_var_type), POINTER                        :: vels
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(section_vals_type), POINTER                   :: tc_section, vc_section
      INTEGER, INTENT(IN)                                :: output_unit, istep

      INTEGER                                            :: i, temp_tol_steps
      LOGICAL                                            :: explicit
      REAL(KIND=dp)                                      :: ext_temp, f_annealing, scale, temp_tol, &
                                                            temploc, tmp_r1, tmp_r2
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: temperatures

! Temperature control

      CALL section_vals_get(tc_section, explicit=explicit)
      IF (explicit) THEN
         CALL section_vals_val_get(tc_section, "TEMP_TOL_STEPS", i_val=temp_tol_steps)
         CALL section_vals_val_get(tc_section, "TEMPERATURE", r_val=ext_temp)
         CALL section_vals_val_get(tc_section, "TEMP_TOL", r_val=temp_tol)
         ALLOCATE (temperatures(SIZE(vels%wrk, 2)))
         ! Computes temperatures
         CALL get_temperatures(vels, particle_set, temperatures, factor=1.0_dp)
         ! Possibly rescale
         IF (istep <= temp_tol_steps) THEN
            DO i = 2, SIZE(vels%wrk, 2) - 1
               temploc = temperatures(i)
               IF (ABS(temploc - ext_temp) > temp_tol) THEN
                  IF (output_unit > 0) THEN
                     tmp_r1 = cp_unit_from_cp2k(temploc, "K")
                     tmp_r2 = cp_unit_from_cp2k(ext_temp, "K")
                     WRITE (output_unit, '(T2,"NEB| Replica Nr.",I5,'// &
                            '"  - Velocity rescaled from: ",F12.6," to: ",F12.6,".")') &
                        i, tmp_r1, tmp_r2

                  END IF
                  scale = SQRT(ext_temp/temploc)
                  vels%wrk(:, i) = scale*vels%wrk(:, i)
               END IF
            END DO
         END IF
         DEALLOCATE (temperatures)
      END IF
      ! Annealing
      CALL section_vals_get(vc_section, explicit=explicit)
      IF (explicit) THEN
         CALL section_vals_val_get(vc_section, "ANNEALING", r_val=f_annealing)
         DO i = 2, SIZE(vels%wrk, 2) - 1
            vels%wrk(:, i) = f_annealing*vels%wrk(:, i)
         END DO
      END IF
   END SUBROUTINE control_vels_a

! **************************************************************************************************
!> \brief Control on velocities - II part
!> \param vels ...
!> \param forces ...
!> \param vc_section ...
!> \author Teodoro Laino 09.2006
! **************************************************************************************************
   SUBROUTINE control_vels_b(vels, forces, vc_section)
      TYPE(neb_var_type), POINTER                        :: vels, forces
      TYPE(section_vals_type), POINTER                   :: vc_section

      INTEGER                                            :: i
      LOGICAL                                            :: explicit, lval
      REAL(KIND=dp)                                      :: factor, norm

! Check the sign of V.dot.F

      CALL section_vals_get(vc_section, explicit=explicit)
      IF (explicit) THEN
         CALL section_vals_val_get(vc_section, "PROJ_VELOCITY_VERLET", l_val=lval)
         IF (lval) THEN
            DO i = 2, SIZE(vels%wrk, 2) - 1
               norm = DOT_PRODUCT(forces%wrk(:, i), forces%wrk(:, i))
               factor = DOT_PRODUCT(vels%wrk(:, i), forces%wrk(:, i))
               IF (factor > 0 .AND. (norm >= EPSILON(0.0_dp))) THEN
                  vels%wrk(:, i) = factor/norm*forces%wrk(:, i)
               ELSE
                  vels%wrk(:, i) = 0.0_dp
               END IF
            END DO
         END IF
         CALL section_vals_val_get(vc_section, "SD_LIKE", l_val=lval)
         IF (lval) THEN
            DO i = 2, SIZE(vels%wrk, 2) - 1
               vels%wrk(:, i) = 0.0_dp
            END DO
         END IF
      END IF
   END SUBROUTINE control_vels_b

! **************************************************************************************************
!> \brief Computes temperatures
!> \param vels ...
!> \param particle_set ...
!> \param temperatures ...
!> \param ekin ...
!> \param factor ...
!> \par History
!>      24.11.2010 rewritten to include core-shell model (MK)
!> \author Teodoro Laino 09.2006
! **************************************************************************************************
   SUBROUTINE get_temperatures(vels, particle_set, temperatures, ekin, factor)

      TYPE(neb_var_type), POINTER                        :: vels
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      REAL(KIND=dp), DIMENSION(:), INTENT(OUT)           :: temperatures
      REAL(KIND=dp), DIMENSION(:), INTENT(OUT), OPTIONAL :: ekin
      REAL(KIND=dp), INTENT(IN), OPTIONAL                :: factor

      INTEGER                                            :: i_rep, iatom, ivar, n_rep, natom, nvar
      REAL(KIND=dp)                                      :: akin, mass, myfactor, temploc
      REAL(KIND=dp), DIMENSION(3)                        :: v

      myfactor = kelvin
      IF (PRESENT(factor)) myfactor = factor
      IF (PRESENT(ekin)) ekin(:) = 0.0_dp
      temperatures(:) = 0.0_dp
      nvar = SIZE(vels%wrk, 1)
      n_rep = SIZE(vels%wrk, 2)
      natom = SIZE(particle_set)
      DO i_rep = 2, n_rep - 1
         akin = 0.0_dp
         IF (vels%in_use == do_band_collective) THEN
            DO ivar = 1, nvar
               akin = akin + 0.5_dp*massunit*vels%wrk(ivar, i_rep)*vels%wrk(ivar, i_rep)
            END DO
         ELSE
            DO iatom = 1, natom
               mass = particle_set(iatom)%atomic_kind%mass
               v(1:3) = get_particle_pos_or_vel(iatom, particle_set, vels%wrk(:, i_rep))
               akin = akin + 0.5_dp*mass*DOT_PRODUCT(v(1:3), v(1:3))
            END DO
            nvar = 3*natom
         END IF
         IF (PRESENT(ekin)) ekin(i_rep) = akin
         temploc = 2.0_dp*akin/REAL(nvar, KIND=dp)
         temperatures(i_rep) = temploc*myfactor
      END DO

   END SUBROUTINE get_temperatures

END MODULE neb_md_utils
